from http.server import BaseHTTPRequestHandler, HTTPServer
import socketserver
import ssl
import json
import os
import re
import winreg
import subprocess
import pymongo
import psutil
from urllib.parse import urlparse, parse_qs
from get_ip_real_address import get_ip_real_address
import time

CERT = os.path.join(os.path.dirname(__file__), 'localhost.crt.pem')  # 公钥
KEY  = os.path.join(os.path.dirname(__file__), 'localhost.key.pem')  # 私钥

PORT_HTTP = 10999
PORT_HTTPS = 10998

def close_specific_consoles():
    current_pid = os.getpid()
    current_process = psutil.Process(current_pid)
    
    # 获取当前Python进程的所有祖先进程
    ancestors = set()
    parent = current_process.parent()
    while parent:
        ancestors.add(parent.pid)
        parent = parent.parent()
    
    for proc in psutil.process_iter(['pid', 'name', 'cmdline', 'create_time']):
        try:
            # 只针对cmd.exe进程
            print(proc.info['name'])
            if 'cmd.exe' in proc.info['name'].lower():
                # 如果这个cmd进程不在当前进程的祖先链中
                if proc.info['pid'] not in ancestors:
                    # 进一步检查命令行参数，确认是独立的cmd窗口
                    cmdline = proc.info['cmdline'] or []
                    cmd_str = ' '.join(cmdline).lower()
                    
                    # 排除正在运行Python的cmd
                    if 'python' not in cmd_str:
                        print(f"关闭独立的CMD进程: PID {proc.info['pid']}")
                        proc.terminate()
                    print(f'Found process: {cmd_str}')
                        
        except (psutil.NoSuchProcess, psutil.AccessDenied):
            continue


def set_user_env_var(name: str, value: str) -> None:
    # 打开当前用户的“环境变量”注册表键
    with winreg.OpenKey(
        winreg.HKEY_CURRENT_USER,
        r"Environment",
        0,
        winreg.KEY_SET_VALUE
    ) as key:
        # 写入/覆盖变量（REG_SZ 即可）
        winreg.SetValueEx(key, name, 0, winreg.REG_SZ, value)

# 添加获取环境变量的函数
def get_user_env_var(name: str) -> str:
    try:
        with winreg.OpenKey(
            winreg.HKEY_CURRENT_USER,
            r"Environment",
            0,
            winreg.KEY_READ
        ) as key:
            value, _ = winreg.QueryValueEx(key, name)
            return value
    except FileNotFoundError:
        return ""

class SimpleHTTPRequestHandler(BaseHTTPRequestHandler):
    def handle_update_mongodb_link(self, post_data):
        """处理更新MongoDB连接链接的请求"""
        data = json.loads(post_data.decode('utf-8'))
        
        # Extract URL field
        url = data.get('url', '')
        if not url:
            self.send_response(400)
            self.end_headers()
            self.wfile.write(b'Missing URL field')
            return
        
        # Parse host and port from URL
        match = re.match(r'tcp://([^:]+):(\d+)', url)
        if not match:
            self.send_response(400)
            self.end_headers()
            self.wfile.write(b'Invalid URL format')
            return
        
        host, port = match.groups()
        
        # Update environment variables
        set_user_env_var('MONGODB_HOST', host)
        set_user_env_var('MONGODB_PORT', port)
        # os.system(f'setx MYMONGODB_HOST "{host}"')
        # os.system(f'setx MYMONGODB_PORT "{port}"')
        
        self.send_response(200)
        self.end_headers()
        self.wfile.write(json.dumps({'result': 'Success'}).encode('utf-8'))

    def handle_get_sku_costs_and_stocks(self):
        """处理获取SKU成本和库存的请求"""
        content_length = int(self.headers['Content-Length'])
        post_data = self.rfile.read(content_length)
        data = json.loads(post_data.decode('utf-8'))
        
        sku_list = data.get('skuList', [])
        if not sku_list:
            self.send_response(400)
            self.end_headers()
            self.wfile.write(b'Missing sku_list field')
            return
        
        results = get_sku_costs_and_stocks(sku_list)
        
        self.send_response(200)
        self.send_header('Content-type', 'application/json')
        self.end_headers()
        self.wfile.write(json.dumps({'results': results}).encode('utf-8'))
        
    def do_POST(self):
        if self.path == '/update_mongodb_link':
            content_length = int(self.headers['Content-Length'])
            post_data = self.rfile.read(content_length)
            self.handle_update_mongodb_link(post_data)
        elif self.path == '/get_sku_costs_and_stocks':
            self.handle_get_sku_costs_and_stocks()
        else:
            self.send_response(404)
            self.end_headers()
            self.wfile.write(b'Not Found')

    # 添加GET方法实现
    def do_GET(self):
        print(f"GET request: {self.path}")
        
        # 获取客户端IP地址
        client_ip = self.client_address[0]
        
        # 如果是本地请求，则允许通过
        if client_ip == "127.0.0.1" or client_ip.startswith("192.168.") or client_ip.startswith("10."):
            pass  # 本地请求直接通过
        else:
            # 检查IP地址是否来自深圳
            location = get_ip_real_address(client_ip)
            if "深圳" not in location:
                self.send_response(403)  # Forbidden
                self.send_header('Content-type', 'application/json')
                self.end_headers()
                error_response = {'error': f'Access denied. Your location ({location}) is not authorized.'}
                self.wfile.write(json.dumps(error_response, ensure_ascii=False).encode('utf-8'))
                return
            
        parsed_url = urlparse(self.path)
        if parsed_url.path == '/update_mongodb_link':
            # 获取当前环境变量
            host = get_user_env_var('MONGODB_HOST')
            port = get_user_env_var('MONGODB_PORT')
            
            # 构造响应数据
            response_data = {
                'url': f'tcp://{host}:{port}' if host and port else ''
            }
            
            # 发送响应
            self.send_response(200)
            self.send_header('Content-type', 'application/json')
            self.end_headers()
            self.wfile.write(json.dumps(response_data).encode('utf-8'))
        elif parsed_url.path == '/fetch_realtime_sells':
            self.start_file_execution('fetch-temu-sells.bat')
        elif parsed_url.path == '/start_daily_crawl':
            self.start_file_execution('auto-fetch-seller-backend-daily.bat')
        elif parsed_url.path == '/relogin_temu_seller':
            self.start_file_execution('auto-relogin.bat')
        elif parsed_url.path == '/get_sku_cost':
            self.handle_get_sku_cost()
        elif parsed_url.path == '/restart_server_task':
            self.start_file_execution('管理员重启计划任务.bat')
        elif parsed_url.path == '/pull_temu_spiders':
            self.start_file_execution('pull_from_gitee.bat')
        elif parsed_url.path.startswith('/get_ip_location'):
            self.handle_get_ip_location()
        elif parsed_url.path == '/close_consoles':
            # 调用关闭特定控制台的函数
            close_specific_consoles()
            # 发送成功响应
            self.send_response(200)
            self.send_header('Content-type', 'application/json')
            self.end_headers()
            self.wfile.write(b'Successfully closed console windows')
        else:
            self.send_response(404)
            self.end_headers()
            self.wfile.write(b'Not Found')
            
    def handle_get_sku_cost(self):
        """处理获取SKU成本的请求"""
        try:
            # 解析URL参数
            parsed_url = urlparse(self.path)
            query_params = parse_qs(parsed_url.query)
            
            # 获取skc和sku参数
            skc = query_params.get('skc', [''])[0]
            sku = query_params.get('sku', [''])[0]
            
            if not skc or not sku:
                self.send_response(400)
                self.send_header('Content-type', 'application/json')
                self.end_headers()
                error_response = {'error': 'Missing required parameters: skc and sku'}
                self.wfile.write(json.dumps(error_response).encode('utf-8'))
                return
            
            # 获取成本信息
            cost_info = get_sku_cost(skc, sku)
            
            if cost_info:
                self.send_response(200)
                self.send_header('Content-type', 'application/json')
                self.end_headers()
                self.wfile.write(json.dumps(cost_info).encode('utf-8'))
            else:
                self.send_response(404)
                self.send_header('Content-type', 'application/json')
                self.end_headers()
                error_response = {'error': f'No cost information found for SKC: {skc}, SKU: {sku}'}
                self.wfile.write(json.dumps(error_response).encode('utf-8'))
        except Exception as e:
            self.send_response(500)
            self.send_header('Content-type', 'application/json')
            self.end_headers()
            error_response = {'error': f'Server error: {str(e)}'}
            self.wfile.write(json.dumps(error_response).encode('utf-8'))
            
    def handle_get_ip_location(self):
        """处理获取IP地址真实位置的请求"""
        try:
            # 解析URL参数
            parsed_url = urlparse(self.path)
            query_params = parse_qs(parsed_url.query)
            
            # 获取ip参数
            ip = query_params.get('ip', [''])[0]
            
            if not ip:
                self.send_response(400)
                self.send_header('Content-type', 'application/json')
                self.end_headers()
                error_response = {'error': 'Missing required parameter: ip'}
                self.wfile.write(json.dumps(error_response).encode('utf-8'))
                return
            
            # 获取IP地址真实位置
            location = get_ip_real_address(ip)
            
            # 发送响应
            self.send_response(200)
            self.send_header('Content-type', 'application/json')
            self.end_headers()
            response_data = {
                'ip': ip,
                'location': location
            }
            self.wfile.write(json.dumps(response_data, ensure_ascii=False).encode('utf-8'))
        except Exception as e:
            self.send_response(500)
            self.send_header('Content-type', 'application/json')
            self.end_headers()
            error_response = {'error': f'Server error: {str(e)}'}
            self.wfile.write(json.dumps(error_response).encode('utf-8'))

    def start_file_execution(self, filename):
        # 获取环境变量 TEMU_SPIDERS_HOME
        spiders_home = os.getenv('TEMU_SPIDERS_HOME')
        if not spiders_home:
            self.send_response(400)
            self.end_headers()
            self.wfile.write(b'Environment variable TEMU_SPIDERS_HOME not set')
            return
        
        # 构造批处理文件路径
        batch_file = os.path.join(spiders_home, filename)
        if not os.path.exists(batch_file):
            self.send_response(400)
            self.end_headers()
            self.wfile.write(f'Batch file not found: {batch_file}'.encode('utf-8'))
            return
        
        # 启动批处理文件并显示窗口
        try:
            print(f"Attempting to start batch file: {batch_file}")
            # 使用绝对路径启动批处理文件
            full_path = os.path.abspath(batch_file)
            print(f"Full path to batch file: {full_path}")
            # 获取批处理文件所在目录
            working_dir = os.path.dirname(full_path)
            print(f"Working directory: {working_dir}")
            # 显式设置窗口显示标志
            startupinfo = subprocess.STARTUPINFO()
            startupinfo.dwFlags |= subprocess.STARTF_USESHOWWINDOW
            # 在批处理文件所在目录中启动进程
            subprocess.Popen(f'start "" "{full_path}"', shell=True, cwd=working_dir)
            print(f"Batch file started: {full_path} in directory: {working_dir}")
            
            # 发送成功响应
            self.send_response(200)
            self.end_headers()
            self.wfile.write(b'Batch file started successfully 5')
        except Exception as e:
            print(f"Failed to start batch file: {e}")
            self.send_response(500)
            self.end_headers()
            self.wfile.write(f'Failed to start batch file: {e}'.encode('utf-8'))
            return

def get_latest_inventory_by_skc_sku(db, skc, sku_code):
    """
    根据SKC和SKU货号获取最新的三个库存数据之和
    
    Args:
        client: pymongo数据库连接对象
        skc: SKC货号
        sku_code: SKU货号
    
    Returns:
        int: 三个库存数据之和，如果未找到返回0
    """
    collection = db['temu_product_sell']
    
    # 查询指定SKC和SKU货号的记录，按时间倒序排列获取最新记录
    record = collection.find_one(
        {
            "skc": skc,
            # "sku_infos": {
            #     "$elemMatch": {
            #         "SKU货号": sku_code
            #     }
            # }
        },
        sort=[("datetime", -1)]  # 按时间降序，获取最新记录
    )
    
    if not record:
        print(f"未找到SKC: {skc}, SKU货号: {sku_code} 的记录")
        return 0
    
    # 在sku_infos数组中查找匹配的SKU
    target_sku = None
    for sku_info in record.get('sku_infos', []):
        if sku_info.get('SKU货号') == sku_code:
            target_sku = sku_info
            break
    
    if not target_sku:
        print(f"在SKC: {skc} 中未找到SKU货号: {sku_code}")
        return 0
    
    # 计算三个库存数据之和
    warehouse_available = target_sku.get('库存数据-仓内可用库存', 0) or 0
    warehouse_unavailable = target_sku.get('库存数据-仓内暂不可用库存', 0) or 0
    shipped = target_sku.get('库存数据-已发货库存', 0) or 0
    
    total_inventory = warehouse_available + warehouse_unavailable + shipped
    
    return total_inventory

def get_latest_inventory_by_skc_sku_bulk(db, skc_sku_pairs: list) -> dict:
    """
    根据SKC和SKU货号获取最新的三个库存数据之和
    
    Args:
        client: pymongo数据库连接对象
        skc: SKC货号
        sku_code: SKU货号
    
    Returns:
        int: 三个库存数据之和，如果未找到返回0
    """
    collection = db['temu_product_sell']
    
    # 查询指定SKC和SKU货号的记录，按时间倒序排列获取最新记录
    records = collection.aggregate([
        {
            '$match': {
                '$or': [
                    {
                        'skc': pair['skc'],
                    } for pair in skc_sku_pairs
                ]
            }
        }, {
            '$sort': {
                'datetime': -1
            }
        }, {
            '$group': {
                '_id': '$skc', 
                'doc': {
                    '$first': '$$ROOT'
                }
            }
        }, {
            '$unwind': '$doc.sku_infos'
        }, {
            '$replaceRoot': {
                'newRoot': '$doc'
            }
        }, {
            '$match': {
                '$or': [
                    {
                        'skc': pair['skc'], 
                        'sku_infos.SKU货号': pair['sku']
                    } for pair in skc_sku_pairs
                ]
            }
        }, {
            '$project': {
                '_id': 0, 
                'skc': 1, 
                'sku': '$sku_infos.SKU货号', 
                '库存数据-仓内可用库存': '$sku_infos.库存数据-仓内可用库存', 
                '库存数据-仓内暂不可用库存': '$sku_infos.库存数据-仓内暂不可用库存', 
                '库存数据-已发货库存': '$sku_infos.库存数据-已发货库存'
            }
        }
    ])

    results = {}
    cursor = list(records)
    for doc in cursor:
        skc = doc['skc']
        sku = doc['sku']
        warehouse_available = doc.get('库存数据-仓内可用库存', 0) or 0
        warehouse_unavailable = doc.get('库存数据-仓内暂不可用库存', 0) or 0
        shipped = doc.get('库存数据-已发货库存', 0) or 0
        total_inventory = warehouse_available + warehouse_unavailable + shipped
        results[f'{skc}_{sku}'] = {
            'skc': skc,
            'sku': sku,
            'inventory': total_inventory
        }

    return results

def get_latest_cost_by_skc_sku(db, skc, sku):
    collection = db["temu_sku_cost"]
    
    # 查询记录
    query = {"skc": int(skc), "SKU货号": sku}
    record = collection.find_one(query)
    
    if not record:
        print(f"No record found for SKC: {skc}, SKU: {sku}")
        return 0
    
    # 获取成本历史中的最后一个成本价
    cost_history = record.get("成本历史", [])
    if not cost_history:
        print(f"No cost history found for SKC: {skc}, SKU: {sku}")
        return 0
    
    latest_cost = cost_history[-1].get("cost")

    return latest_cost

def get_latest_cost_by_skc_sku_bulk(db, skc_sku_pairs: list) -> dict:
    """
    批量获取指定 skc 和 sku 的最新成本记录（使用 MongoDB 聚合管道优化）
    :param skc_sku_pairs: 包含 skc 和 sku 的字典列表，格式为 [{"skc": "SKC001", "sku": "SKU001"}, ...]
    :return: 返回一个字典，键为 skc_sku 组合，值为对应的最新成本记录
    """
    if not skc_sku_pairs:
        return {}
    
    collection = db["temu_sku_cost"]

    # 构建聚合管道
    pipeline = [
        {
            "$match": {
                "$or": [
                    {
                        "skc": int(pair["skc"]),
                        "SKU货号": pair["sku"]
                    }
                    for pair in skc_sku_pairs
                ]
            }
        },
        {
            "$sort": {
                "datetime": -1  # 按时间降序排序
            }
        },
        {
            "$group": {
                "_id": {
                    "skc": "$skc",
                    "sku": "$SKU货号"
                },
                "latest_record": {
                    "$first": "$$ROOT"  # 获取每组的第一条记录（即最新记录）
                }
            }
        }
    ]

    # 执行聚合查询
    cursor = list(collection.aggregate(pipeline))

    # 处理结果
    results = {}
    for doc in cursor:
        skc = doc["_id"]["skc"]
        sku = doc["_id"]["sku"]
        latest_record = doc["latest_record"]

        cost_history = latest_record.get("成本历史", [])
        if not cost_history:
            print(f"No cost history found for SKC: {skc}, SKU: {sku}")
            latest_cost = 0
        else:
            latest_cost = cost_history[-1].get("cost")
        
        results[f"{skc}_{sku}"] = {
            "skc": skc,
            "sku": sku,
            "cost": latest_cost,
        }

    return results

def get_sku_costs_and_stocks(sku_list):
    """
    从MongoDB获取指定SKC和SKU的成本价
    :param skc: 产品SKC编号
    :param sku: SKU货号
    :return: 包含skc、sku和cost的字典，如果未找到则返回None
    """
    try:
        # 获取MongoDB连接信息
        host = get_user_env_var('MONGODB_HOST')
        port = get_user_env_var('MONGODB_PORT')
        username = get_user_env_var('MONGODB_USERNAME')
        password = get_user_env_var('MONGODB_PASSWORD')
        dbname = get_user_env_var('MONGODB_DBNAME') or "temu_seller_db"  # 默认数据库名
        
        if not host or not port:
            print("MongoDB connection info not found in environment variables")
            return None
        
        # 构建MongoDB连接URI
        if username and password:
            uri = f"mongodb://{username}:{password}@{host}:{port}/{dbname}"
        else:
            uri = f"mongodb://{host}:{port}/{dbname}"
            
        results = []
        # 连接MongoDB
        with pymongo.MongoClient(uri) as client:
            db = client[dbname]

            costs = get_latest_cost_by_skc_sku_bulk(db, sku_list)

            inventories = get_latest_inventory_by_skc_sku_bulk(db, sku_list)

            for skc_sku in inventories:
                cost = costs.get(skc_sku, {}).get('cost', 0)
                inventories[skc_sku].update({'cost': cost})
            # for item in sku_list:
            #     skc = item.get('skc', '')
            #     sku = item.get('sku', '')
            #     time1 = time.time()
            #     latest_cost = get_latest_cost_by_skc_sku(db, skc, sku)
            #     time2 = time.time()
            #     print(f"Time taken to retrieve cost for SKC: {skc}, SKU: {sku}: {time2 - time1} seconds")
            #     time1 = time2

            #     latest_inventory = get_latest_inventory_by_skc_sku(db, skc, sku)
            #     print(f"Time taken to retrieve inventory for SKC: {skc}, SKU: {sku}: {time.time() - time1} seconds")
                
            #     results.append({
            #         "skc": skc,
            #         "sku": sku,
            #         "cost": latest_cost,
            #         "inventory": latest_inventory
            #     })

            results = list(inventories.values())
        
        return results
    except Exception as e:
        print(f"Error retrieving SKU cost: {e}")
        return None

def get_sku_cost(skc, sku):
    """
    从MongoDB获取指定SKC和SKU的成本价
    :param skc: 产品SKC编号
    :param sku: SKU货号
    :return: 包含skc、sku和cost的字典，如果未找到则返回None
    """
    try:
        # 获取MongoDB连接信息
        host = get_user_env_var('MONGODB_HOST')
        port = get_user_env_var('MONGODB_PORT')
        username = get_user_env_var('MONGODB_USERNAME')
        password = get_user_env_var('MONGODB_PASSWORD')
        dbname = get_user_env_var('MONGODB_DBNAME') or "temu_seller_db"  # 默认数据库名
        
        if not host or not port:
            print("MongoDB connection info not found in environment variables")
            return None
        
        # 构建MongoDB连接URI
        if username and password:
            uri = f"mongodb://{username}:{password}@{host}:{port}/{dbname}"
        else:
            uri = f"mongodb://{host}:{port}/{dbname}"
            
        # 连接MongoDB
        client = pymongo.MongoClient(uri)
        db = client[dbname]

        latest_cost = get_latest_cost_by_skc_sku(db, skc, sku)

        latest_inventory = get_latest_inventory_by_skc_sku(db, skc, sku)
        
        return {
            "skc": skc,
            "sku": sku,
            "cost": latest_cost,
            "inventory": latest_inventory
        }
    except Exception as e:
        print(f"Error retrieving SKU cost: {e}")
        return None

def run(server_class=HTTPServer, handler_class=SimpleHTTPRequestHandler, port=PORT_HTTP):
    server_address = ('', port)
    httpd = server_class(server_address, handler_class)
    print(f'Starting httpd on port {port}...')
    httpd.serve_forever()

def run_with_ssl(handler_class=SimpleHTTPRequestHandler, port=PORT_HTTPS):

    PORT = port  # HTTPS 默认端口是 443，这里用 4443 避免权限问题

    # Create SSL context using recommended approach
    context = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
    context.load_cert_chain(certfile=CERT, keyfile=KEY)

    with socketserver.TCPServer(("", PORT), handler_class) as httpd:
        # 包装 socket 以支持 SSL
        httpd.socket = context.wrap_socket(httpd.socket, server_side=True)
        print(f"Serving HTTPS on port {PORT}")
        httpd.serve_forever()

if __name__ == '__main__':
    # 如果存在SSL证书文件，则使用HTTPS启动服务
    if os.path.exists(CERT) and os.path.exists(KEY):
        run_with_ssl()
    else:
        print("SSL certificate not found. Starting HTTP server.")
        run()